import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Performs general matrix multiplication (GEMM): C = A @ B
    
    Args:
        A (torch.Tensor): Input matrix of shape (M, K)
        B (torch.Tensor): Input matrix of shape (K, N)
    
    Returns:
        torch.Tensor: Output matrix of shape (M, N)
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """
        Performs matrix multiplication C = A @ B.

        Args:
            A: Input tensor of shape (M, K).
            B: Input tensor of shape (K, N).

        Returns:
            C: Output tensor of shape (M, N).
        """
        return torch.matmul(A, B)


# Default dimensions
M = 1024
K = 1024
N = 1024


def get_inputs():
    """
    Generate random input tensors for the model.
    
    Returns:
        List[torch.Tensor]: List containing tensors A and B
    """
    A = torch.randn(M, K, dtype=torch.float16)
    B = torch.randn(K, N, dtype=torch.float16)
    return [A, B]


def get_init_inputs():
    """
    Generate tensors required for model initialization.
    
    Returns:
        List: Empty list as no special initialization is needed
    """
    return []  # No special initialization inputs needed
